'''
Created on Aug 19, 2009

@author: mkiyer
'''
import sys
import h5py
import numpy as np
from optparse import OptionParser
import os

def read_bed_file(fhd):
    for line in fhd:
        fields = line.strip().split('\t')
        yield (fields[0], int(fields[1]), int(fields[2]))

def plot_scatter_hist(x, y, outfile):
    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.ticker import NullFormatter
    
    nullfmt = NullFormatter()         # no labels
    # definitions for the axes
    left, width = 0.1, 0.65
    bottom, height = 0.1, 0.65
    bottom_h = left_h = left+width+0.02
    
    rect_scatter = [left, bottom, width, height]
    rect_histx = [left, bottom_h, width, 0.2]
    rect_histy = [left_h, bottom, 0.2, height]
    
    # start with a rectangular Figure
    plt.figure(1, figsize=(8,8))
    
    axScatter = plt.axes(rect_scatter)
    axHistx = plt.axes(rect_histx)
    axHisty = plt.axes(rect_histy)
    
    # no labels
    axHistx.xaxis.set_major_formatter(nullfmt)
    axHisty.yaxis.set_major_formatter(nullfmt)
    
    # the scatter plot:
    axScatter.scatter(x, y)
    
    # now determine nice limits by hand:
    binwidth = 0.25
    xymax = np.max( [np.max(np.fabs(x)), np.max(np.fabs(y))] )
    lim = ( int(xymax/binwidth) + 1) * binwidth
    
    axScatter.set_xlim( (-lim, lim) )
    axScatter.set_ylim( (-lim, lim) )
    
    bins = np.arange(-lim, lim + binwidth, binwidth)
    axHistx.hist(x, bins=bins)
    axHisty.hist(y, bins=bins, orientation='horizontal')
    
    axHistx.set_xlim( axScatter.get_xlim() )
    axHisty.set_ylim( axScatter.get_ylim() )    
    plt.savefig(outfile)
    plt.close()

def plot_overlapping_hist(x, y, outfile, x_name="group1", y_name="group2"):
    import pylab as P
    # histogram
    P.figure()
    # the histogram of the data
    nbins = 100
    #n, bins, patches = P.hist([x,y], nbins, alpha=1)
    n, bins, patches = P.hist(x, bins=nbins, facecolor='g', alpha=0.5)        
    n2, bins2, patches2 = P.hist(y, bins=bins, facecolor='r', alpha=0.5)
    #n, bins, patches = plt.hist(exoncov, bins=50, log=True, facecolor='r', alpha=0.5)
    #n2, bins2, patches2 = plt.hist(covarray, bins, log=True, facecolor='g', alpha=0.5)
    #, normed=1, facecolor='g', alpha=0.75)
    #n, bins, patches = plt.hist([lander_exoncov, refflat_exoncov, covarray], bins=50, range=(0, 50), log=True)
    P.legend([patches[0], patches2[0]], [x_name, y_name])

    P.xlabel('normalized coverage')
    P.ylabel('bin count')
    P.title('H3K36 coverage comparison')
    P.grid()
    #plt.axis([40, 160, 0, 0.03])
    P.savefig(outfile)
    P.close()

def get_interval_coverage(intervals, h5file):
    h5data = h5py.File(h5file, 'r')
    for chrom, start, end in intervals:
        if chrom in h5data:
            dset = h5data[chrom]
            cov = 1000 * (np.sum(dset[start:end]) / float(end - start))
            yield cov
        else:
            yield 0.0    

if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options] sample_file")
    optionparser.add_option("--hdf5", dest="h5",
                            help="h5 file containing coverage data")
    optionparser.add_option("-x", dest="x",
                            help="BED file containing intervals")
    optionparser.add_option("-y", dest="y",
                            help="BED file containing intervals")
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")
    (options, args) = optionparser.parse_args()
   
    if options.outfile == None:
        optionparser.error("no output file specified..")
        #options.outfile = os.path.splitext(os.path.basename(interval_file))[0] + ".png"

    xintervals = list(read_bed_file(open(options.x)))
    yintervals = list(read_bed_file(open(options.y)))

    print 'getting cov for', options.x
    xcov = [c for c in get_interval_coverage(xintervals, options.h5)]
    print 'getting cov for', options.y
    ycov = [c for c in get_interval_coverage(yintervals, options.h5)]
    plot_overlapping_hist(xcov, ycov, options.outfile, 
                          x_name=os.path.splitext(os.path.basename(options.x))[0],
                          y_name=os.path.splitext(os.path.basename(options.y))[0])

